# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

import datetime

import coin.strategy.mm.fastfeature.feature_feed_converter as ffcv
import coin.strategy.mm.fastfeature.feed_manager as fmgr
from coin.strategy.mm.base.bookmap import BookMap
import coin.strategy.mm.tool.archive_base as abase
import coin.strategy.mm.fastfeature.sim.lm_sim_executor as lmsimexec
import coin.strategy.mm.fastfeature.sim.lm_strategy_logic as lmlogic
import coin.strategy.mm.fastfeature.sim.lm_agg_strategy_logic as lmagglogic
import coin.strategy.executor.aggressive.util as autil
import coin.strategy.executor.passive.util as putil
from coin.strategy.hk_lm.lm_agg_strategy import LmAggStrategyConfig
from coin.strategy.hk_hard_arb.util import load_clean_json

import coin.proto.archive_flag_pb2 as afpb
import pyfastfeature.feature.simulator as pysim
from coin.base.datetime_util import to_timestamp_int

import os
from coin.support.latency.latency_model import LatencyModel


class LmSimStrategy(object):
  StrategyLogicType = None
  ExecutorType = None
  PostOnly = False

  def __init__(self,
               *,
               logger,
               model_profile,
               trade_product,
               latency_ms,
               delay_accept_until_book,
               start_time,
               end_time):
    self._start_time = to_timestamp_int(start_time)
    self._end_time = to_timestamp_int(end_time)
    self._logger = logger

    self._feed_products = tuple(model_profile.product_map.keys())
    self._feed_converter = model_profile.converter_factory.create_feed_converter(
        ffcv.FeatureFeedCacheConverter)

    self._config = model_profile.execution_params
    self._config['leverage'] = 10
    self._execution_params = type(self).ExecutorConfigConverter(trade_product.symbol, self._config)

    self._trade_product = trade_product

    self._trade_subreq = model_profile.sub_key
    self._feed_subreqs = model_profile.feed_subreqs
    self._model = model_profile.model
    self._book = BookMap([trade_product], logger=self._logger)
    self._executor = None
    self._exit = False

    self._feed_manager = fmgr.FeedManager(self.on_book,
                                          self.on_trade, {},
                                          feedspec=model_profile.converter_factory.feedspec)

    self._simulator = pysim.Simulator(list(model_profile.product_map.values()),
                                      latency_ns=int(latency_ms * 1e6),
                                      delay_accept_until_book=delay_accept_until_book)

    self._feed_converter.register_symbols(self._model, self._simulator)
    self._trade_pindex = self._feed_converter.get_index(trade_product)

    order_gateway = lmsimexec.LmSimOrderGateway(self._logger,
                                                self._simulator,
                                                trade_product,
                                                self._trade_pindex,
                                                post_only=self.PostOnly)

    if self.StrategyLogicType == lmagglogic.LmAggLogic:
      self._lm_logic = self.StrategyLogicType(config=self._config,
                                              trade_product=trade_product,
                                              logger=self._logger)
    else:
      self._lm_logic = self.StrategyLogicType(config=self._config,
                                              trade_product=trade_product,
                                              trade_pindex=self._trade_pindex,
                                              logger=self._logger)

    self._executor = self.ExecutorType(trade_product,
                                       self._execution_params,
                                       order_gateway,
                                       self._logger)

  def get_misc_info(self):
    misc_info = self._lm_logic.get_misc_info()
    if misc_info is None:
      misc_info = {}
    misc_info['initial_pos'] = self._executor.og.get_initial_position(self._trade_product)
    return misc_info

  def on_book_reset(self, book_builder_name, book_builder):
    self._feed_manager.set_book_reset(self._feed_products, book_builder_name, book_builder)

  def try_fire(self, timestamp):
    self._lm_logic.try_fire_and_manage(timestamp=timestamp,
                                       executor=self._executor,
                                       feed_converter=self._feed_converter,
                                       model=self._model,
                                       book=self._book.get_book(self._trade_product),
                                       dry_run=False)

  def on_trade(self, product, trade):
    if self._start_time < trade.timestamp < self._end_time:
      if not self._feed_converter.update_trade(product, trade, self._model):
        return
      self._feed_converter.update_trade(product, trade, self._simulator)
      self.try_fire(trade.timestamp)

  def on_book(self, product, book):
    if self._start_time < book.timestamp < self._end_time:
      if not self._feed_converter.update_books(product, book, [self._model, self._simulator]):
        return
      if product == self._trade_product:
        self._book.set_book(product, book)
      self.try_fire(book.timestamp)
    else:
      self._executor.og.cancel_all()


class LmAggSimStrategy(LmSimStrategy):
  StrategyLogicType = lmagglogic.LmAggLogic
  ExecutorType = lmsimexec.LmSimAggExecutor
  ExecutorConfigConverter = autil.convert_executor_config_from_dict


class LmPassThickBookSimStrategy(LmSimStrategy):
  StrategyLogicType = lmlogic.LmPassThickBookLogic
  ExecutorType = lmsimexec.LmSimMMExecutor
  ExecutorConfigConverter = putil.convert_executor_config_from_dict
  PostOnly = True


class LmPassThinBookSimStrategy(LmSimStrategy):
  StrategyLogicType = lmlogic.LmPassThinBookLogic
  ExecutorType = lmsimexec.LmSimMMExecutor
  ExecutorConfigConverter = putil.convert_executor_config_from_dict
  PostOnly = True


class LmPassPricerSimStrategy(LmSimStrategy):
  StrategyLogicType = lmlogic.LmPassPricerLogic
  ExecutorType = lmsimexec.LmSimMMExecutor
  ExecutorConfigConverter = putil.convert_executor_config_from_dict
  PostOnly = True


class LmSimMultiStrategy(object):
  def __init__(self,
               strat_config,
               start_time,
               end_time,
               latency_ms,
               is_agg,
               logger,
               use_predicted_latency,
               initial_position=None,
               latency_multiplier=None,
               delay_accept_until_book=None):
    assert isinstance(strat_config, LmAggStrategyConfig)

    if use_predicted_latency:
      model = LatencyModel()
      path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                          '../../../../support/latency/20200103.model')
      model.load(path)
      try:
        latency_ms = int(
            model.predict({
                'exchange': strat_config._trade_subreq.exchange,
                'market_type': strat_config._trade_subreq.market_type
            }) / 1e6)
        if latency_multiplier is None:
          latency_multiplier = 1.0
        latency_ms *= latency_multiplier or 1
        print('predicted latency * latency_multiplier: ',
              latency_ms,
              'ms for ',
              strat_config._trade_subreq.exchange,
              strat_config._trade_subreq.market_type)
      except Exception:
        pass

    self.substrats = {}
    for name, profile in strat_config.model_profile.items():
      thin_or_thickbook = profile.execution_params.get("thin_or_thickbook_mode", "thick")
      if is_agg:
        self.substrats[name] = LmAggSimStrategy(model_profile=profile,
                                                trade_product=profile.product,
                                                latency_ms=latency_ms,
                                                delay_accept_until_book=delay_accept_until_book,
                                                start_time=start_time,
                                                end_time=end_time,
                                                logger=logger)
      elif thin_or_thickbook == "pricer":
        self.substrats[name] = LmPassPricerSimStrategy(
            model_profile=profile,
            trade_product=profile.product,
            latency_ms=latency_ms,
            delay_accept_until_book=delay_accept_until_book,
            start_time=start_time,
            end_time=end_time,
            logger=logger)
      elif thin_or_thickbook == "thin":
        self.substrats[name] = LmPassThinBookSimStrategy(
            model_profile=profile,
            trade_product=profile.product,
            latency_ms=latency_ms,
            delay_accept_until_book=delay_accept_until_book,
            start_time=start_time,
            end_time=end_time,
            logger=logger)
      elif thin_or_thickbook == "thick":
        self.substrats[name] = LmPassThickBookSimStrategy(
            model_profile=profile,
            trade_product=profile.product,
            latency_ms=latency_ms,
            delay_accept_until_book=delay_accept_until_book,
            start_time=start_time,
            end_time=end_time,
            logger=logger)
      else:
        raise ValueError(thin_or_thickbook)
      if initial_position is not None:
        self.substrats[name]._executor.og.set_initial_position(profile.product,
                                                               float(initial_position))

  def on_book_reset(self, book_builder_name, book_builder):
    for substrat in self.substrats.values():
      substrat.on_book_reset(book_builder_name, book_builder)

  def get_products(self):
    products = []
    for substrat in self.substrats.values():
      products.extend(list(substrat._feed_products))
    return list(set(products))

  def get_result(self):
    ret_dict = {}
    for name, substrat in self.substrats.items():
      ret_dict[name] = (substrat._trade_product,
                        substrat._simulator.GetHistoricalResponse(substrat._trade_pindex),
                        substrat.get_misc_info())
    return ret_dict


def get_strategy_config(*,
                        config_filename,
                        trading_date,
                        time_range,
                        threshold_bps=None,
                        cancel_threshold_bps=None,
                        taker_threshold_bps=None,
                        edge_allowance_bps=None,
                        sticky_bp=None,
                        max_posting_period_sec=None,
                        order_update_period_sec=None,
                        lot_size_mult=None,
                        preserve_stack=None,
                        stack=None,
                        neutral=None,
                        pick_symbol=None,
                        **sink_kwargs):
  add_hour, hours = time_range.split("-")
  add_hour = float(add_hour)
  hours = float(hours) - add_hour
  if trading_date is None:
    start_time = None
  else:
    start_time = trading_date + datetime.timedelta(hours=add_hour)
  clean_json = load_clean_json(config_filename)

  if pick_symbol is not None and len(pick_symbol) > 0:
    clean_json['products'] = {
        key: value for key, value in clean_json['products'].items() if key.find(pick_symbol) >= 0
    }

  strat_config = LmAggStrategyConfig(clean_json, start_time)

  for model_prof_value in strat_config.model_profile.values():
    for set_column in ['threshold_bps', 'edge_allowance_bps']:
      if eval(set_column) is not None:
        model_prof_value.execution_params[set_column] = float(eval(set_column))

      if neutral:
        reserve = 0.5 * \
            (model_prof_value.execution_params['max_pos']
             + model_prof_value.execution_params['min_pos'])
        model_prof_value.execution_params['max_pos'] -= reserve
        model_prof_value.execution_params['min_pos'] -= reserve

      if preserve_stack:
        assert neutral is True
        assert stack is None
        model_prof_value.execution_params['max_pos'] *= float(lot_size_mult)
        model_prof_value.execution_params['min_pos'] *= float(lot_size_mult)

      if stack is not None:
        model_prof_value.execution_params['max_pos'] = float(
            model_prof_value.execution_params['lot_size']) * (float(stack) / 2)
        model_prof_value.execution_params['min_pos'] = - \
            float(model_prof_value.execution_params['lot_size']) * (float(stack) / 2)

  if config_filename.find("pass") >= 0:
    for model_prof_value in strat_config.model_profile.values():
      for set_column in ['cancel_threshold_bps', 'taker_threshold_bps', 'sticky_bp']:
        if eval(set_column) is not None:
          model_prof_value.execution_params[set_column] = float(eval(set_column))

      if lot_size_mult is not None:
        model_prof_value.execution_params['lot_size'] *= float(lot_size_mult)
      if order_update_period_sec is not None:
        model_prof_value.execution_params['order_update_period_sec'] = float(
            order_update_period_sec)
        model_prof_value.execution_params['min_posting_period_sec'] = float(order_update_period_sec)
  else:
    for model_prof_value in strat_config.model_profile.values():
      if max_posting_period_sec is not None:
        model_prof_value.execution_params['max_posting_period'] = float(
            max_posting_period_sec) * 1e9
        model_prof_value.execution_params['max_posting_period_sec'] = float(max_posting_period_sec)

    for model_prof_value in strat_config.model_profile.values():
      if lot_size_mult is not None:
        model_prof_value.execution_params['lot_size'] *= float(lot_size_mult)

  return strat_config


def dump_lm_sim(logger, **sink_kwargs):
  add_hour, hours = sink_kwargs['time_range'].split("-")
  add_hour = float(add_hour)
  hours = float(hours) - add_hour
  start_time = sink_kwargs['trading_date'] + datetime.timedelta(hours=add_hour)
  end_time = start_time + datetime.timedelta(hours=hours)

  strat_config = get_strategy_config(**sink_kwargs)

  dumpstrategy = LmSimMultiStrategy(strat_config,
                                    start_time=start_time,
                                    end_time=end_time,
                                    latency_ms=sink_kwargs['latency_ms'],
                                    is_agg=(sink_kwargs['config_filename'].find("pass") < 0),
                                    logger=logger,
                                    use_predicted_latency=sink_kwargs['use_predicted_latency'],
                                    latency_multiplier=sink_kwargs['latency_multiplier'],
                                    initial_position=sink_kwargs['initial_position'],
                                    delay_accept_until_book=sink_kwargs['delay_accept_until_book'])

  baseflags = afpb.BaseFlagProto(trading_date=sink_kwargs['trading_date'].strftime("%Y%m%d"),
                                 time_range=sink_kwargs['time_range'])

  fsflags = abase.get_feed_subscription_flags()

  abase.run_from_archive_base_with_feed_sub(baseflags=baseflags,
                                            feed_sub_flags_proto=fsflags,
                                            feed_sub_request=strat_config._feed_subreq,
                                            on_book_reset=dumpstrategy.on_book_reset)

  return dumpstrategy.get_result()
